/**
 * 
 */
package edu.umd.clip.lm.model;

import edu.umd.clip.lm.factors.*;
import edu.umd.clip.lm.util.*;

import java.io.*;
import java.util.*;

import com.sleepycat.bind.tuple.*;
/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class ProbTree implements Cloneable, Externalizable {
	protected BinaryPrefix prefix;

	private static final double PRUNING_THRESHOLD = 1E-9;
	static final double MIN_TREE_PROB = 1E-10;
	//long tuple;
	//BinaryPrefix wordPrefix;
	double scale;
	DecoderCompactProbTree tree;
	//HashSet<BinaryPrefix> cuts;
	boolean readonly = false;
	
	/**
	 * Constructs a probability tree of s single tuple (used for <s>)
	 * @param tuple
	 */
	public ProbTree() {
		this(new BinaryPrefix(0));
	}
	
	public ProbTree(BinaryPrefix prefix) {
		scale = 1.0;
		tree = DecoderCompactProbTree.getSingleNodeTree();
		this.prefix = prefix;		
	}
	/**
	 * @param overtFactors
	 * @param prefix
	 * @param logScale
	 */
	private ProbTree(BinaryPrefix prefix, double scale) {
		this.prefix = prefix;
		this.scale = scale;
	}

	public ProbTree compact() {
		ProbTree t = new ProbTree(prefix, scale);
		t.tree = tree == null ? null : tree.compact();
		return t;
	}
	
	public ProbTree(DecoderCompactProbTree t) {
		this();
		tree = t;
	}

	public ProbTree(TupleInput in) throws IndexOutOfBoundsException {
		this();
		readTuple(in);
	}

	public ProbTree(ObjectInput in) throws IOException, ClassNotFoundException {
		this();
		readExternal(in);
	}

	public static ProbTree makeProbTree(OnDiskCompactProbTree leafy) {
		long start = Timing.nanoTime();
		try {
			DecoderCompactProbTree tree = (DecoderCompactProbTree) DecoderCompactProbTree.getEmptyTree().clone();
			Experiment exp = Experiment.getInstance();
			FactorTupleDescription desc = exp.getTupleDescription();
			
			int compactFactors[] = leafy.getCompactHiddenFactors();
			float probabilities[] = leafy.getProbabilities();
			for(int i=0; i<leafy.getSize(); ++i) {
				long hiddenFactors = desc.unpackHiddenFactorsFromInt(compactFactors[i]);
				tree.addProbability(exp.getHiddenPrefix(hiddenFactors), probabilities[i]);
			}
			
			ProbTree probTree = new ProbTree(tree);
			probTree.scale = leafy.getScale();
			return probTree;
		} finally {
			long end = Timing.nanoTime();
			totalCompactifyingTime += end - start;
		}
	}
	
	public static long totalLeafyMergeTime = 0;
	public static long totalLeafyMergeCount = 0;
	public static long totalCompactMergeTime = 0;
	public static long totalCompactMergeCount = 0;
	public static long totalCloningTime = 0;
	public static long totalMergingCount = 0;
	public static long totalCompactifyingTime = 0;
	
	/*
	public void addTree(ProbTree other) {
		this.tree.scale(this.scale / other.scale);
		this.scale = other.scale;
		
		if (other.tree.isCompactProbTree()) {
			long start = nanoTime(); 
			this.tree.getCompactProbTree().addTree(other.tree.getCompactProbTreeReadOnly());
			long end = nanoTime();
			if (TIMING) {
				totalCompactMergeTime += end - start;
				totalCompactMergeCount++;
			}
		} else {
			OnDiskCompactProbTree leafy;
			if (other.tree instanceof CompactProbTreeProxy) {
				leafy = (OnDiskCompactProbTree) ((CompactProbTreeProxy) other.tree).getImpl();
			} else {
				leafy = (OnDiskCompactProbTree) other.tree;
			}
			long start = nanoTime(); 
			for(OnDiskCompactProbTree.PathProb path : leafy.getPaths()) {
				this.tree.addProbability(path.path, path.prob);
			}
			long end = nanoTime();
			if (TIMING) {
				totalLeafyMergeTime += end - start;
				totalLeafyMergeCount++;
			}
		}
	}
	
	public static ProbTree mergeProbTrees(Collection<ProbTree> trees) {
		if (trees.size() == 0) return null;
		
		if (TIMING) 
			++totalMergingCount;
		DecoderCompactProbTree newCompactTree;
		{
			long start = nanoTime();
			newCompactTree = (DecoderCompactProbTree) DecoderCompactProbTree.getEmptyTree().clone();
			long end = nanoTime();
			if (TIMING) {
				totalCloningTime += end - start;
			}
		}
		ArrayList<DecoderCompactProbTree> compactTrees = new ArrayList<DecoderCompactProbTree>(trees.size());
		double scale = 0.0;
		for(ProbTree t : trees) {
			if (t.scale > scale) scale = t.scale;
		}
		if (scale == 0.0) {
			return null;
		}
		long leafyMergeCount = 0;
		long leafyMergeTime = 0;
		
		for(ProbTree t : trees) {
			t.tree.scale(t.scale / scale);
			t.scale = scale;
			if (t.tree.isCompactProbTree()) {
				compactTrees.add(t.tree.getCompactProbTreeReadOnly());
			} else {
				OnDiskCompactProbTree leafy;
				if (t.tree instanceof CompactProbTreeProxy) {
					leafy = (OnDiskCompactProbTree) ((CompactProbTreeProxy) t.tree).getImpl();
				} else {
					leafy = (OnDiskCompactProbTree) t.tree;
				}
				long start = System.nanoTime(); 
				for(OnDiskCompactProbTree.PathProb path : leafy.getPaths()) {
					newCompactTree.addProbability(path.path, path.prob);
				}
				long end = System.nanoTime();
				leafyMergeTime += end - start;
				leafyMergeCount++;
			}
		}
		// presumably overt factors are the same for all trees
		long overtFactors = trees.iterator().next().tuple;
		
		{
			long start = nanoTime();
			if (compactTrees.size() == 2) {
				newCompactTree.addTwoTrees(compactTrees.get(0), compactTrees.get(1));
			} else if (compactTrees.size() > 0) {
				newCompactTree.addTrees(compactTrees);
			}
			long end = nanoTime();
			if (TIMING) {
				totalCompactMergeTime += end - start;
				totalCompactMergeCount += compactTrees.size();
			}
		}
		if (TIMING) {
			totalLeafyMergeTime += leafyMergeTime;
			totalLeafyMergeCount += leafyMergeCount;
		}
		
		double totalProb = newCompactTree.getTotalProb();
		assert(totalProb >= -1e-6);
		if (totalProb <= 0) return null;

		AbstractCompactProbTree compactTree = newCompactTree;
		newCompactTree.scale(1.0/totalProb);
		
		long start = nanoTime();
		//newCompactTree = newCompactTree.prune(PRUNING_THRESHOLD);
		if (newCompactTree == null) return null;
		
		AbstractCompactProbTree t = new CompactProbTreeProxy(newCompactTree);
		//compactTree = t.compact();
		compactTree = t;
		long end = nanoTime();
		if (TIMING)
			totalCompactifyingTime += end - start;
		
		if (compactTree == null) return null;
		
		ProbTree tree = new ProbTree(overtFactors, new BinaryPrefix(0), totalProb*scale);
		tree.tree = compactTree;
		return tree;
	}
	*/
	
	/**
	 * tests whether how the cuttingPrefix will cut the tree
	 * @param cuttingPrefix
	 * @return TRUE if the tree is a subtree of the cuttingPrefix, 
	 *         FALSE if there's no intersection between the tree and cuttingPrefix,
	 *         null if the cuttingPrefix can cut the tree (the intersection may still be empty because of pruning)  
	 */
	public Boolean testCut(BinaryPrefix cuttingPrefix) {
		if (cuttingPrefix.isPrefixOf(prefix)) {
			return Boolean.TRUE;
		}
		if (!prefix.isPrefixOf(cuttingPrefix)) {
			return Boolean.FALSE;
		}
		if (tree.findNode(cuttingPrefix.getSuffix(prefix.getSize())) == DecoderCompactProbTree.NO_LINK) {
			return Boolean.FALSE;
		}
		/*
		if (cuts != null) {
			for(BinaryPrefix p : cuts) {
				if (p.isPrefixOf(cuttingPrefix)) return Boolean.FALSE;
			}
		}
		*/
		return null;
	}
	/**
	 * Chops a subtree identified by cuttingPrefix off the ProbTree
	 * @param cuttingPrefix
	 * @return the subtree cut, null if nothing has been cut, returns this if all leaves are cut
	 */
	public ProbTree cut(BinaryPrefix cuttingPrefix) {
		DecoderCompactProbTree myTree = tree.getCompactProbTree();
		if (readonly) {
			tree = myTree.copy();
			myTree = tree.getCompactProbTree();
			readonly = false;
		}
		DecoderCompactProbTree cutTree = myTree.cut(cuttingPrefix.getSuffix(prefix.getSize()));
		if (cutTree == null || cutTree.getTotalProb() < ProbMath.MIN_FLOAT_PRECISION) return null;
		if (myTree == cutTree || myTree.getTotalProb() < ProbMath.MIN_FLOAT_PRECISION) {
			//the prefix cut off all leaves from the tree 
			this.tree = cutTree;
			this.prefix = cuttingPrefix;
			return this;
		}
		ProbTree subTree = new ProbTree(cuttingPrefix, scale);
		subTree.tree = cutTree;
		/*
		if (cuts != null) {
			Iterator<BinaryPrefix> it = cuts.iterator();
			while(it.hasNext()) {
				BinaryPrefix p = it.next();
				if (cuttingPrefix.isPrefixOf(p)) {
					it.remove();
					if (subTree.cuts == null) subTree.cuts = new HashSet<BinaryPrefix>();
					subTree.cuts.add(p);
				}
			}
		} else {
			cuts = new HashSet<BinaryPrefix>();
		}
		cuts.add(cuttingPrefix);
		*/
		return subTree;
	}
	
	/*
	public double getProbability(int packedHiddenFactors) {
		return tree.getProbability(packedHiddenFactors) * scale;
	}
	
	public void setProbability(int packedHiddenFactors, double prob) {
		tree.setProbability(packedHiddenFactors, prob / scale);
	}
	
	public void addProbability(int packedHiddenFactors, double prob) {
		tree.addProbability(packedHiddenFactors, prob / scale);
	}
	*/
	/* (non-Javadoc)
	 * @see java.lang.Object#clone()
	 */
	@SuppressWarnings("unchecked")
	@Override
	public Object clone() {
		ProbTree copy = new ProbTree(prefix, scale);
		copy.tree = (DecoderCompactProbTree) tree.clone();
		/*
		if (cuts != null) {
			copy.cuts = (HashSet<BinaryPrefix>) cuts.clone();
		}
		*/
		return copy;
	}

	@SuppressWarnings("unchecked")
	public ProbTree copy() {
		ProbTree copy = new ProbTree(prefix, scale);
		//copy.tree = tree.getCompactProbTree().copy();
		copy.tree = tree;
		readonly = true;
		copy.readonly = true;
		/*
		if (cuts != null) {
			copy.cuts = (HashSet<BinaryPrefix>) cuts.clone();
		}
		*/
		return copy;
	}

	/**
	 * @return the prefix
	 */
	public BinaryPrefix getPrefix() {
		return prefix;
	}

	/**
	 * @return the cuts
	 */
	/*
	public HashSet<BinaryPrefix> getCuts() {
		return cuts;
	}
*/
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.questions.AbstractProbTree#getRoot()
	 */
	public BinaryPrefix getRoot() {
		return prefix;
	}

	/**
	 * @return the logScale
	 */
	public double getScale() {
		return scale;
	}
	
	public double getTreeProb() {
		return tree.getTotalProb();
	}

	public double getTotalProb() {
		return getTreeProb() * scale;
	}
	
	public void scale(double scale) {
		this.scale *= scale;
	}

	@Override
	public int hashCode() {
		final int prime = 31;
		int result = 1;
		long temp;
		temp = Double.doubleToLongBits(scale);
		result = prime * result + (int) (temp ^ (temp >>> 32));
		result = prime * result + ((tree == null) ? 0 : tree.hashCode());
		return result;
	}

	@Override
	public boolean equals(Object obj) {
		if (this == obj)
			return true;
		if (obj == null)
			return false;
		if (!(obj instanceof ProbTree))
			return false;
		ProbTree other = (ProbTree) obj;
		if (!prefix.equals(other.prefix))
			return false;
		if (scale != other.scale)
			return false;
		
		return tree.equals(other.tree);
	}

	/*
	public AbstractCompactProbTree getTree() {
		return tree;
	}
	
	public void setTree(AbstractCompactProbTree tree) {
		this.tree = tree;
	}
	*/
	
	public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException {
		Experiment experiment = Experiment.getInstance();
		
		scale = in.readDouble();
		tree = new DecoderCompactProbTree(in);//.getCompactProbTreeReadOnly().compact();
		short numCuts = in.readShort();
		/*
		cuts = null;
		if (numCuts > 0) {
			cuts = new HashSet<BinaryPrefix>(numCuts);
			for(short i=0; i<numCuts; ++i) {
				cuts.add(experiment.getCachedPrefix(new BinaryPrefix(in)));
			}
		}
		*/
	}

	public void writeExternal(ObjectOutput out) throws IOException {
		out.writeDouble(scale);
		tree.writeExternal(out);
		/*
		if (cuts != null) {
			out.writeShort((short) cuts.size());
			for(BinaryPrefix p : cuts) {
				p.writeExternal(out);
			}
		} else {
			out.writeShort((short) 0);
		}
		*/
	}

	public void readTuple(TupleInput in) throws IndexOutOfBoundsException {
		Experiment experiment = Experiment.getInstance();
		
		scale = in.readDouble();
		tree = new DecoderCompactProbTree(in);//.getCompactProbTreeReadOnly().compact();
		short numCuts = in.readShort();
		/*
		cuts = null;
		if (numCuts > 0) {
			cuts = new HashSet<BinaryPrefix>(numCuts);
			for(short i=0; i<numCuts; ++i) {
				cuts.add(experiment.getCachedPrefix(new BinaryPrefix(in)));
			}
		}
		*/
	}

	public void writeTuple(TupleOutput out) {
		out.writeDouble(scale);
		tree.writeTuple(out);
		/*
		if (cuts != null) {
			out.writeShort((short) cuts.size());
			for(BinaryPrefix p : cuts) {
				p.writeTuple(out);
			}
		} else {
			out.writeShort((short) 0);
		}
		*/
	}

	public ProbTree getProbTree() {
		return this;
	}
	
}
